
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import pandas as pd
import numpy as np
import seaborn as sns
from sklearn.feature_selection import SelectKBest
import random
from django.conf import settings
import os

NAME_LENGTH = 20


def get_randon_abs_file_path():
    letter = '0123456789abcdefghiklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
    filename = ''
    for i in range(NAME_LENGTH):
        filename += letter[random.randint(0, len(letter) - 1)]
    return os.path.join('pics', filename + '.png')


def get_heatmap_path(x_data: pd.DataFrame):
    crr_matrix = x_data.corr()
    f, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(crr_matrix, annot=True, cmap='YlGnBu')
    ax.set_title('Features Correlation')
    filename = get_randon_abs_file_path()
    f.savefig(os.path.join(settings.MEDIA_ROOT, filename), dpi=100, bbox_inches='tight')
    return filename


def get_feature_imp_path(x_data: pd.DataFrame, y_data: pd.DataFrame):
    f, ax = plt.subplots(figsize=(10, 8))
    ax.set_title('Features Importance')
    sclt = SelectKBest(k='all')
    sclt.fit(x_data, y_data)
    score = sclt.scores_
    features = x_data.columns
    named_score = zip(features, score)
    sorted_named_score = sorted(named_score, key=lambda z: z[1], reverse=True)
    sorted_score = [each[1] for each in sorted_named_score]
    sorted_name = [each[0] for each in sorted_named_score]
    y_pos = np.arange(len(features))
    plt.barh(y_pos, sorted_score, height=0.7, align='center', color='#AAAAAA', tick_label=sorted_name)
    plt.xlabel('Feature Score')
    plt.ylabel('Feature Name')
    for score, pos in zip(sorted_score, y_pos):
        plt.text(score + 20, pos, '%.1f' % score, ha='center', va='bottom', fontsize=8)
    filename = get_randon_abs_file_path()
    f.savefig(os.path.join(settings.MEDIA_ROOT, filename), dpi=100, bbox_inches='tight')
    return filename


def get_pca_3d(X_pca, Y):
    x = X_pca[:, 0]
    y = X_pca[:, 1]
    z = Y
    plt.style.use('ggplot')
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(x, y, z)
    ax.set_title('PCA')
    ax.set_zlabel('Z')  # 坐标轴
    ax.set_ylabel('Y')
    ax.set_xlabel('X')
    filename = get_randon_abs_file_path()
    fig.savefig(os.path.join(settings.MEDIA_ROOT, filename), dpi=100)
    return filename


if __name__ == '__main__':
    print(get_randon_abs_file_path())
